import os

import matplotlib.pyplot as plt
import numpy as np
import torch
from torch.utils.data import DataLoader

from CRF import DenseCRF
from SwinUnet import SwinUNet
from train import MedicalDataset


def test(model, dataloader, device):
    model.eval()
    # 创建 DenseCRF 实例
    crf = DenseCRF(iter_max=100, pos_w=130, pos_xy_std=1, bi_w=10, bi_xy_std=80, bi_rgb_std=13)

    with torch.no_grad():
        for i, test_X in enumerate(dataloader):
            if i >= 10:  # 只显示前10张图片
                break

            # 将输入图像移到设备并进行预测
            test_X = test_X.float().to(device)
            pred_Y = model(test_X)

            # 如果模型输出是列表，选择最后一个输出
            if isinstance(pred_Y, list):
                pred_Y = pred_Y[-1]

            # 转换为NumPy格式用于显示
            test_X_np = test_X[0].permute(1, 2, 0).cpu().numpy()
            pred_Y_np = pred_Y[0].squeeze(0).cpu().numpy()
            pred_Y_np = (pred_Y_np > 0.5).astype(np.float32)  # 二值化

            # 反归一化处理以显示原始图像
            mean = np.array([0.485, 0.456, 0.406])
            std = np.array([0.229, 0.224, 0.225])
            test_X_display = (test_X_np * std) + mean
            test_X_display = np.clip(test_X_display, 0, 1)  # 限制数值范围以避免显示错误

            # 应用DenseCRF
            probs = np.stack([1 - pred_Y_np, pred_Y_np], axis=0)  # 假设是二分类问题
            image = (test_X_display * 255).astype(np.uint8)  # 转换回原始图像格式用于CRF

            refined_pred = crf(image, probs)
            refined_pred = np.argmax(refined_pred, axis=0)  # 获取最有可能的类别

            # 创建画布，并添加子图
            plt.figure(figsize=(12, 4))

            # 显示输入图像
            plt.subplot(1, 3, 1)
            plt.imshow(test_X_display)
            plt.title("Input Image")
            plt.xticks([])
            plt.yticks([])
            plt.axis('off')

            # 显示原始预测的分割掩码
            plt.subplot(1, 3, 2)
            plt.imshow(pred_Y_np, cmap='gray')
            plt.title("Original Predicted Mask")
            plt.xticks([])
            plt.yticks([])
            plt.axis('off')

            # 显示DenseCRF优化后的分割掩码
            plt.subplot(1, 3, 3)
            plt.imshow(refined_pred, cmap='gray')
            plt.title("Refined Predicted Mask by DenseCRF")
            plt.xticks([])
            plt.yticks([])
            plt.axis('off')

            plt.tight_layout()
            plt.show()

            print(f"Processed image {i + 1}/10")


if __name__ == '__main__':

    # 创建测试数据集
    test_dataset = MedicalDataset(
        root_dir='/Volumes/For_Mac/dateset/Pulmonary_X_ray_and_masks',
        is_train=False,
        image_size=224
    )

    # 创建数据加载器
    test_loader = DataLoader(
        test_dataset,
        batch_size=4,
        shuffle=False,
        num_workers=4
    )

    # 初始化模型
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = SwinUNet(224, 224, 3, 96, 1, 3, 4)

    # 加载训练好的模型参数
    model_path = '../pt_file/4X-SwinUnet-66.pt'  # 确保这是你保存模型的正确路径
    if os.path.exists(model_path):
        model.load_state_dict(torch.load(model_path, map_location=device))
        print(f"Model loaded from {model_path}")
    else:
        print(f"Model file {model_path} does not exist.")
        exit()

    # 开始测试
    test(model, test_loader, device)
